// 全局变量
let model = null;
let webcam = null;
let trainingData = {
    class1: { images: [], labels: [], name: '类别A' },
    class2: { images: [], labels: [], name: '类别B' },
    class3: { images: [], labels: [], name: '类别C' }
};
let trainingChart = null;
let isPredicting = false;

// 初始化
document.addEventListener('DOMContentLoaded', function() {
    initializeApp();
});

function initializeApp() {
    // 初始化标签页
    showTab('home');
    
    // 检查摄像头支持
    if (!navigator.mediaDevices || !navigator.mediaDevices.getUserMedia) {
        alert('您的浏览器不支持摄像头功能');
    }
    
    // 初始化图表
    initTrainingChart();
}

// 标签页切换
function showTab(tabName) {
    // 隐藏所有标签页
    document.querySelectorAll('.tab-content').forEach(tab => {
        tab.style.display = 'none';
    });
    
    // 显示选中的标签页
    document.getElementById(tabName + '-tab').style.display = 'block';
    
    // 特殊处理
    if (tabName === 'train') {
        startTrainingCameras();
    } else if (tabName === 'predict') {
        startPredictCamera();
    }
}

// 启动摄像头
async function startCamera() {
    try {
        const video = document.getElementById('preview-video');
        const stream = await navigator.mediaDevices.getUserMedia({ 
            video: { width: 320, height: 240 } 
        });
        video.srcObject = stream;
    } catch (error) {
        console.error('摄像头启动失败:', error);
        alert('无法启动摄像头，请检查权限设置');
    }
}

// 启动训练页面的摄像头
async function startTrainingCameras() {
    const cameras = ['class1-video', 'class2-video', 'class3-video'];
    
    for (const cameraId of cameras) {
        try {
            const video = document.getElementById(cameraId);
            if (!video.srcObject) {
                const stream = await navigator.mediaDevices.getUserMedia({ 
                    video: { width: 320, height: 240 } 
                });
                video.srcObject = stream;
            }
        } catch (error) {
            console.error(`摄像头 ${cameraId} 启动失败:`, error);
        }
    }
}

// 启动预测页面的摄像头
async function startPredictCamera() {
    try {
        const video = document.getElementById('predict-video');
        if (!video.srcObject) {
            const stream = await navigator.mediaDevices.getUserMedia({ 
                video: { width: 640, height: 480 } 
            });
            video.srcObject = stream;
        }
    } catch (error) {
        console.error('预测摄像头启动失败:', error);
    }
}

// 拍照
function captureImage(classNum) {
    const video = document.getElementById(`class${classNum}-video`);
    const canvas = document.createElement('canvas');
    canvas.width = 224;
    canvas.height = 224;
    const ctx = canvas.getContext('2d');
    
    // 绘制并调整图片大小
    ctx.drawImage(video, 0, 0, 224, 224);
    
    // 转换为张量并存储
    const imageData = ctx.getImageData(0, 0, 224, 224);
    const tensor = tf.browser.fromPixels(imageData);
    
    trainingData[`class${classNum}`].images.push(tensor);
    trainingData[`class${classNum}`].labels.push(classNum - 1);
    
    // 更新UI
    updateImageGrid(classNum);
    updateCount(classNum);
}

// 上传图片
function uploadImages(event, classNum) {
    const files = event.target.files;
    
    Array.from(files).forEach(file => {
        const reader = new FileReader();
        reader.onload = function(e) {
            const img = new Image();
            img.onload = function() {
                const canvas = document.createElement('canvas');
                canvas.width = 224;
                canvas.height = 224;
                const ctx = canvas.getContext('2d');
                ctx.drawImage(img, 0, 0, 224, 224);
                
                const imageData = ctx.getImageData(0, 0, 224, 224);
                const tensor = tf.browser.fromPixels(imageData);
                
                trainingData[`class${classNum}`].images.push(tensor);
                trainingData[`class${classNum}`].labels.push(classNum - 1);
                
                updateImageGrid(classNum);
                updateCount(classNum);
            };
            img.src = e.target.result;
        };
        reader.readAsDataURL(file);
    });
    
    // 清空input，允许重复上传相同文件
    event.target.value = '';
}

// 更新图片网格
function updateImageGrid(classNum) {
    const container = document.getElementById(`class${classNum}-images`);
    container.innerHTML = '';
    
    const images = trainingData[`class${classNum}`].images;
    images.forEach((tensor, index) => {
        const div = document.createElement('div');
        div.className = 'image-item';
        
        const img = document.createElement('img');
        tf.browser.toPixels(tensor, img);
        
        const removeBtn = document.createElement('button');
        removeBtn.className = 'remove-btn';
        removeBtn.innerHTML = '×';
        removeBtn.onclick = () => removeImage(classNum, index);
        
        div.appendChild(img);
        div.appendChild(removeBtn);
        container.appendChild(div);
    });
}

// 移除图片
function removeImage(classNum, index) {
    trainingData[`class${classNum}`].images.splice(index, 1);
    trainingData[`class${classNum}`].labels.splice(index, 1);
    
    updateImageGrid(classNum);
    updateCount(classNum);
}

// 更新计数
function updateCount(classNum) {
    const count = trainingData[`class${classNum}`].images.length;
    document.getElementById(`class${classNum}-count`).textContent = count;
}

// 初始化训练图表
function initTrainingChart() {
    const ctx = document.getElementById('training-chart').getContext('2d');
    trainingChart = new Chart(ctx, {
        type: 'line',
        data: {
            labels: [],
            datasets: [{
                label: 'Loss',
                data: [],
                borderColor: 'rgb(255, 99, 132)',
                backgroundColor: 'rgba(255, 99, 132, 0.1)',
                yAxisID: 'y'
            }, {
                label: 'Accuracy',
                data: [],
                borderColor: 'rgb(54, 162, 235)',
                backgroundColor: 'rgba(54, 162, 235, 0.1)',
                yAxisID: 'y1'
            }]
        },
        options: {
            responsive: true,
            interaction: {
                mode: 'index',
                intersect: false,
            },
            scales: {
                x: {
                    display: true,
                    title: {
                        display: true,
                        text: 'Epoch'
                    }
                },
                y: {
                    type: 'linear',
                    display: true,
                    position: 'left',
                    title: {
                        display: true,
                        text: 'Loss'
                    }
                },
                y1: {
                    type: 'linear',
                    display: true,
                    position: 'right',
                    title: {
                        display: true,
                        text: 'Accuracy'
                    },
                    grid: {
                        drawOnChartArea: false,
                    },
                }
            }
        }
    });
}

// 创建模型
function createModel() {
    const model = tf.sequential({
        layers: [
            tf.layers.conv2d({
                inputShape: [224, 224, 3],
                filters: 8,
                kernelSize: 5,
                strides: 1,
                activation: 'relu',
                kernelInitializer: 'varianceScaling'
            }),
            tf.layers.maxPooling2d({ poolSize: [2, 2], strides: [2, 2] }),
            tf.layers.conv2d({
                filters: 16,
                kernelSize: 3,
                activation: 'relu'
            }),
            tf.layers.maxPooling2d({ poolSize: [2, 2], strides: [2, 2] }),
            tf.layers.flatten(),
            tf.layers.dense({
                units: 32,
                activation: 'relu'
            }),
            tf.layers.dense({
                units: 3,
                activation: 'softmax'
            })
        ]
    });
    
    return model;
}

// 开始训练
async function startTraining() {
    // 检查数据量
    const totalImages = trainingData.class1.images.length + 
                       trainingData.class2.images.length + 
                       trainingData.class3.images.length;
    
    if (totalImages < 60) {
        alert('请为每个类别至少收集20张图片');
        return;
    }
    
    // 禁用训练按钮
    const trainButton = document.querySelector('button[onclick="startTraining()"]');
    trainButton.disabled = true;
    trainButton.innerHTML = '<span class="loading"></span> 训练中...';
    
    // 更新进度条
    const progressBar = document.getElementById('training-progress');
    progressBar.style.width = '0%';
    progressBar.textContent = '准备数据...';
    
    // 准备训练数据
    const allImages = [
        ...trainingData.class1.images,
        ...trainingData.class2.images,
        ...trainingData.class3.images
    ];
    
    const allLabels = [
        ...trainingData.class1.labels,
        ...trainingData.class2.labels,
        ...trainingData.class3.labels
    ];
    
    const xs = tf.stack(allImages);
    const ys = tf.oneHot(tf.tensor1d(allLabels, 'int32'), 3);
    
    // 创建和编译模型
    model = createModel();
    const learningRate = parseFloat(document.getElementById('learning-rate').value);
    const optimizer = tf.train.adam(learningRate);
    
    model.compile({
        optimizer: optimizer,
        loss: 'categoricalCrossentropy',
        metrics: ['accuracy']
    });
    
    // 训练参数
    const epochs = parseInt(document.getElementById('epochs').value);
    const batchSize = Math.floor(allImages.length / 10);
    
    // 训练模型
    const history = await model.fit(xs, ys, {
        epochs: epochs,
        batchSize: batchSize,
        validationSplit: 0.2,
        callbacks: {
            onEpochEnd: async (epoch, logs) => {
                // 更新进度条
                const progress = ((epoch + 1) / epochs) * 100;
                progressBar.style.width = progress + '%';
                progressBar.textContent = `第 ${epoch + 1}/${epochs} 轮`;
                
                // 更新图表
                trainingChart.data.labels.push(epoch + 1);
                trainingChart.data.datasets[0].data.push(logs.loss);
                trainingChart.data.datasets[1].data.push(logs.acc);
                trainingChart.update();
                
                // 更新UI
                await tf.nextFrame();
            },
            onTrainEnd: async () => {
                progressBar.style.width = '100%';
                progressBar.textContent = '训练完成！';
                progressBar.classList.remove('progress-bar-animated');
                progressBar.classList.add('bg-success');
                
                // 恢复按钮
                trainButton.disabled = false;
                trainButton.innerHTML = '<i class="fas fa-play"></i> 开始训练';
                
                alert('模型训练完成！现在可以切换到预测页面进行测试。');
            }
        }
    });
    
    // 清理内存
    xs.dispose();
    ys.dispose();
}

// 开始预测
async function startPrediction() {
    if (!model) {
        alert('请先训练模型！');
        return;
    }
    
    isPredicting = true;
    const video = document.getElementById('predict-video');
    const resultsContainer = document.getElementById('prediction-results');
    
    // 获取类别名称
    const classNames = [
        document.getElementById('class1-name').value || '类别A',
        document.getElementById('class2-name').value || '类别B',
        document.getElementById('class3-name').value || '类别C'
    ];
    
    function predict() {
        if (!isPredicting) return;
        
        const canvas = document.createElement('canvas');
        canvas.width = 224;
        canvas.height = 224;
        const ctx = canvas.getContext('2d');
        ctx.drawImage(video, 0, 0, 224, 224);
        
        const imageData = ctx.getImageData(0, 0, 224, 224);
        const tensor = tf.browser.fromPixels(imageData).expandDims(0);
        
        model.predict(tensor).data().then(predictions => {
            // 更新预测结果
            resultsContainer.innerHTML = '';
            
            predictions.forEach((confidence, index) => {
                const percentage = (confidence * 100).toFixed(1);
                
                const item = document.createElement('div');
                item.className = 'prediction-item';
                item.innerHTML = `
                    <div class="prediction-label">${classNames[index]}</div>
                    <div class="progress">
                        <div class="progress-bar ${getProgressBarClass(index)}" 
                             style="width: ${percentage}%">
                            ${percentage}%
                        </div>
                    </div>
                `;
                resultsContainer.appendChild(item);
            });
            
            tensor.dispose();
            
            if (isPredicting) {
                requestAnimationFrame(predict);
            }
        });
    }
    
    predict();
}

// 停止预测
function stopPrediction() {
    isPredicting = false;
}

// 获取进度条颜色类
function getProgressBarClass(index) {
    const classes = ['bg-primary', 'bg-success', 'bg-warning', 'bg-danger', 'bg-info'];
    return classes[index % classes.length];
}

// 清除所有数据
function clearAllData() {
    if (confirm('确定要清除所有收集的数据吗？')) {
        trainingData.class1.images = [];
        trainingData.class1.labels = [];
        trainingData.class2.images = [];
        trainingData.class2.labels = [];
        trainingData.class3.images = [];
        trainingData.class3.labels = [];
        
        updateImageGrid(1);
        updateImageGrid(2);
        updateImageGrid(3);
        updateCount(1);
        updateCount(2);
        updateCount(3);
        
        // 重置图表
        trainingChart.data.labels = [];
        trainingChart.data.datasets[0].data = [];
        trainingChart.data.datasets[1].data = [];
        trainingChart.update();
        
        // 重置进度条
        const progressBar = document.getElementById('training-progress');
        progressBar.style.width = '0%';
        progressBar.textContent = '准备就绪';
        progressBar.classList.remove('bg-success');
    }
}

// 导出模型
async function exportModel() {
    if (!model) {
        alert('请先训练模型！');
        return;
    }
    
    try {
        // 保存模型
        const saveResult = await model.save('downloads://teachable-machine-model');
        
        // 保存类别名称
        const classNames = {
            class1: document.getElementById('class1-name').value || '类别A',
            class2: document.getElementById('class2-name').value || '类别B',
            class3: document.getElementById('class3-name').value || '类别C'
        };
        
        const blob = new Blob([JSON.stringify(classNames)], { type: 'application/json' });
        const url = URL.createObjectURL(blob);
        const a = document.createElement('a');
        a.href = url;
        a.download = 'class-names.json';
        a.click();
        URL.revokeObjectURL(url);
        
        alert('模型导出成功！请保存生成的文件。');
    } catch (error) {
        console.error('导出失败:', error);
        alert('模型导出失败: ' + error.message);
    }
}

// 导入模型
async function importModel(event) {
    const file = event.target.files[0];
    if (!file) return;
    
    try {
        const reader = new FileReader();
        reader.onload = async function(e) {
            const modelJson = JSON.parse(e.target.result);
            
            // 加载模型
            model = await tf.loadLayersModel(tf.io.browserFiles([file]));
            
            alert('模型导入成功！现在可以开始预测了。');
        };
        reader.readAsText(file);
    } catch (error) {
        console.error('导入失败:', error);
        alert('模型导入失败: ' + error.message);
    }
    
    // 清空input
    event.target.value = '';
}

// 页面卸载时清理资源
window.addEventListener('beforeunload', function() {
    if (model) {
        model.dispose();
    }
    
    // 停止所有摄像头
    document.querySelectorAll('video').forEach(video => {
        if (video.srcObject) {
            video.srcObject.getTracks().forEach(track => track.stop());
        }
    });
});
